import os
import sys

# 添加项目根目录到Python路径
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))

from config_reader import get_langsmith_config, get_ollama_host
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.llms import Ollama

# 从配置文件读取LangSmith配置
langsmith_config = get_langsmith_config()

# 设置环境变量
os.environ["LANGCHAIN_TRACING_V2"] = langsmith_config['tracing_v2']
os.environ["LANGCHAIN_ENDPOINT"] = langsmith_config['endpoint']
os.environ["LANGCHAIN_PROJECT"] = langsmith_config['project']
os.environ["LANGCHAIN_API_KEY"] = langsmith_config['api_key']

# 准备测试数据
documents = [
    Document(page_content="狗是伟大的伴侣，以其忠诚和友好而闻名。", metadata={"source": "哺乳动物宠物文档"}),
    Document(page_content="猫是独立的宠物，通常喜欢自己的空间。", metadata={"source": "哺乳动物宠物文档"}),
    Document(page_content="金鱼是初学者的流行宠物，需要相对简单的护理。", metadata={"source": "鱼类宠物文档"}),
    Document(page_content="鹦鹉是聪明的鸟类，能够模仿人类的语言。", metadata={"source": "鸟类宠物文档"}),
    Document(page_content="兔子是社交动物，需要足够的空间跳跃。", metadata={"source": "哺乳动物宠物文档"})
]

# 使用本地 Ollama 模型
ollama_host = get_ollama_host()
model = Ollama(model="llama3.2", base_url=ollama_host)

# 实例化一个向量数空间（使用 Ollama 嵌入）
vector_store = Chroma.from_documents(documents, embedding=OllamaEmbeddings(
    model="mxbai-embed-large",
    base_url=ollama_host
))

# 相似度查询: 返回相似的分数， 分数越低相似度越高
print(vector_store.similarity_search_with_score('咖啡猫'))
